import cv2
import numpy as np


class SparseOpticalFlow(object):
    def __init__(self) -> None:
        super().__init__()
        self.feature_params = dict(maxCorners=500,
                                   qualityLevel=0.01,
                                   minDistance=50,
                                   blockSize=7)

        criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03)
        self.lk_params = dict(winSize=(15, 15), maxLevel=2, criteria=criteria)

    def _detect_good_point(self, img, det_boxes):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        mask_image = np.ones_like(gray) * 255
        img_h, img_w, _ = img.shape

        # 如果使用跟踪框，第1位是ID，此处需要去除
        if len(det_boxes) > 0 and det_boxes.shape[1] >= 4:
            det_boxes = det_boxes[:, 1:]

        for det_box in det_boxes:
            x1, y1, x2, y2 = det_box
            # 矩形框四周向外扩充10，去除边界的影响
            x1 = max(int(x1.item() * img_w) - 10, 0)
            y1 = max(int(y1.item() * img_h) - 10, 0)
            x2 = min(int(x2.item() * img_w) + 10, img_w - 1)
            y2 = min(int(y2.item() * img_h) + 10, img_h - 1)
            mask_image[y1:y2, x1:x2] = 0
        points = cv2.goodFeaturesToTrack(gray,
                                         mask=mask_image,
                                         **self.feature_params)
        return points

    @staticmethod
    def _filter_points(prev_points, curr_points, status):
        # step 1: 基于OpenCV API，保留status为1的points
        prev_points = prev_points[status == 1]
        curr_points = curr_points[status == 1]
        # step 2: 保留低于1个标准差的flow，去除异常点
        distances = np.sqrt(
            np.sum(np.square(prev_points - curr_points), axis=1))
        mean = np.mean(distances)
        std = np.std(distances)
        keep = np.abs(distances - mean) < std
        prev_points = prev_points[keep]
        curr_points = curr_points[keep]
        return prev_points, curr_points

    def optical_flow(self, prev_img, curr_img, prev_det_boxes=None):
        prev_points = self._detect_good_point(prev_img, prev_det_boxes)
        prev_gray = cv2.cvtColor(prev_img, cv2.COLOR_BGR2GRAY)
        curr_gray = cv2.cvtColor(curr_img, cv2.COLOR_BGR2GRAY)
        curr_points, status, err = cv2.calcOpticalFlowPyrLK(
            prev_gray, curr_gray, prev_points, None, **self.lk_params)
        prev_points, curr_points = self._filter_points(prev_points,
                                                       curr_points, status)
        # 使用最小二乘法计算Homography
        flag = np.ones(shape=(prev_points.shape[0], 1),
                       dtype=prev_points.dtype)
        A = np.hstack([prev_points, flag])
        y = np.hstack([curr_points, flag])
        H = np.matmul(np.linalg.inv(np.matmul(A.T, A)), np.matmul(A.T, y))
        H = H.T
        return H
